"""
工具函数模块，提供各种辅助功能
"""
import os
import time
import logging
import torch
from typing import Dict, Any, List, Optional

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

def get_rocm_info() -> Dict[str, Any]:
    """
    获取ROCm相关信息
    
    返回:
        包含ROCm信息的字典
    """
    info = {
        "rocm_available": False,
        "device_count": 0,
        "devices": [],
        "visible_devices": os.environ.get("HIP_VISIBLE_DEVICES", "未设置")
    }
    
    if torch.cuda.is_available():
        info["rocm_available"] = True
        info["device_count"] = torch.cuda.device_count()
        
        for i in range(info["device_count"]):
            device_info = {
                "index": i,
                "name": torch.cuda.get_device_name(i),
                "memory_total": torch.cuda.get_device_properties(i).total_memory / (1024 ** 3),  # GB
                "memory_free": torch.cuda.memory_reserved(i) / (1024 ** 3)  # GB
            }
            info["devices"].append(device_info)
    
    return info

def measure_inference_time(func):
    """
    测量推理时间的装饰器
    
    参数:
        func: 要测量的函数
        
    返回:
        包装后的函数
    """
    def wrapper(*args, **kwargs):
        start_time = time.time()
        result = func(*args, **kwargs)
        end_time = time.time()
        
        # 添加推理时间到结果中
        if isinstance(result, dict):
            if "metadata" not in result:
                result["metadata"] = {}
            result["metadata"]["inference_time"] = end_time - start_time
        
        return result
    
    return wrapper

def format_size(size_bytes: int) -> str:
    """
    格式化字节大小为人类可读格式
    
    参数:
        size_bytes: 字节大小
        
    返回:
        格式化后的字符串
    """
    if size_bytes < 1024:
        return f"{size_bytes} B"
    elif size_bytes < 1024 ** 2:
        return f"{size_bytes / 1024:.2f} KB"
    elif size_bytes < 1024 ** 3:
        return f"{size_bytes / (1024 ** 2):.2f} MB"
    else:
        return f"{size_bytes / (1024 ** 3):.2f} GB"